{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c7a7e67",
   "metadata": {},
   "source": [
    "# Tutorial 03: weak imposition of Dirichlet BCs by a Lagrange multiplier (nonlinear problem)\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{align*}\n",
    "&\\min_{u} \\int_\\Omega \\left\\{ (1 + u^2)\\ |\\nabla u|^2 - u \\right\\} dx,\\\\\n",
    "&\\text{s.t. } u = g\\text{ on }\\Gamma = \\partial \\Omega\n",
    "\\end{align*}$$\n",
    "where $\\Omega$ is a ball in 2D.\n",
    "\n",
    "The optimality conditions result in the following nonlinear problem\n",
    "\n",
    "$$\\begin{align*}\n",
    "&\\int_\\Omega (1+u^2)\\ \\nabla u \\cdot \\nabla v \\ dx + \\int_\\Omega u \\ |\\nabla u|^2 v \\ dx = \\int_\\Omega v \\ dx\\\\\n",
    "&\\text{s.t. } u = g\\text{ on }\\Gamma = \\partial \\Omega\n",
    "\\end{align*}$$\n",
    "\n",
    "\n",
    "We compare the following two cases:\n",
    "* **strong imposition of Dirichlet BCs**:\n",
    "the corresponding weak formulation is\n",
    "$$\n",
    "\\text{find } u \\in V_g \\text{ s.t. } \\int_\\Omega (1+u^2)\\ \\nabla u \\cdot \\nabla v \\ dx + \\int_\\Omega u \\ |\\nabla u|^2 v \\ dx = \\int_\\Omega v \\ dx, \\quad \\forall v \\in V_0\\\\\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V_g = \\{v \\in H^1(\\Omega): v|_\\Gamma = g\\},\\\\\n",
    "V_0 = \\{v \\in H^1(\\Omega): v|_\\Gamma = 0\\}.\\\\\n",
    "$$\n",
    "* **weak imposition of Dirichlet BCs**: this requires an introduction of a multiplier $\\lambda$ which is restricted to $\\Gamma$, and solves\n",
    "\n",
    "$$\n",
    "\\text{find } u, \\lambda \\in V \\times M \\text{ s.t. }\\\\\n",
    "\\begin{align*}\n",
    "&\\int_\\Omega (1+u^2)\\ \\nabla u \\cdot \\nabla v \\ dx + \\int_\\Omega u \\ |\\nabla u|^2 v \\ dx & &+ \\int_\\Gamma \\lambda v \\ dx & &= \\int_\\Omega v \\ dx, & \\forall v \\in V,\\\\\n",
    "&\\int_\\Gamma u \\mu \\ ds & & & &= \\int_\\Gamma g \\mu \\ ds, & \\forall \\mu \\in M\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where\n",
    "$$\n",
    "V = H^1(\\Omega),\\\\\n",
    "M = L^{2}(\\Gamma).\\\\\n",
    "$$\n",
    "\n",
    "This example is a prototypical case of problems containing subdomain/boundary restricted variables (the Lagrange multiplier, in this case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74cf97d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e708bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import gmsh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "921087c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3515090d",
   "metadata": {},
   "source": [
    "### Geometrical parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0deb0ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 3\n",
    "mesh_size = 1. / 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be09ac4d",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d0345ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.initialize()\n",
    "gmsh.model.add(\"mesh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d2b91f",
   "metadata": {},
   "outputs": [],
   "source": [
    "p0 = gmsh.model.geo.addPoint(0.0, 0.0, 0.0, mesh_size)\n",
    "p1 = gmsh.model.geo.addPoint(0.0, +r, 0.0, mesh_size)\n",
    "p2 = gmsh.model.geo.addPoint(0.0, -r, 0.0, mesh_size)\n",
    "c0 = gmsh.model.geo.addCircleArc(p1, p0, p2)\n",
    "c1 = gmsh.model.geo.addCircleArc(p2, p0, p1)\n",
    "boundary = gmsh.model.geo.addCurveLoop([c0, c1])\n",
    "domain = gmsh.model.geo.addPlaneSurface([boundary])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f13b244",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.model.geo.synchronize()\n",
    "gmsh.model.addPhysicalGroup(1, [c0, c1], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [boundary], 0)\n",
    "gmsh.model.mesh.generate(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a20b6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, subdomains, boundaries = dolfinx.io.gmshio.model_to_mesh(\n",
    "    gmsh.model, comm=mpi4py.MPI.COMM_WORLD, rank=0, gdim=2)\n",
    "gmsh.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d329a5e5-42b5-4f11-8484-9fee49b0794a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab47b49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "facets_Gamma = boundaries.indices[boundaries.values == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9335b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea947631",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d74bb465",
   "metadata": {},
   "source": [
    "### Weak imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eccbf35a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function space\n",
    "V = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "M = V.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24db5fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_V = np.arange(0, V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts)\n",
    "dofs_M_Gamma = dolfinx.fem.locate_dofs_topological(M, boundaries.dim, facets_Gamma)\n",
    "restriction_V = multiphenicsx.fem.DofMapRestriction(V.dofmap, dofs_V)\n",
    "restriction_M_Gamma = multiphenicsx.fem.DofMapRestriction(M.dofmap, dofs_M_Gamma)\n",
    "restriction = [restriction_V, restriction_M_Gamma]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00a9f11b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions, as well as solution\n",
    "(du, dl) = (ufl.TrialFunction(V), ufl.TrialFunction(M))\n",
    "(u, l) = (dolfinx.fem.Function(V), dolfinx.fem.Function(M))\n",
    "(v, m) = (ufl.TestFunction(V), ufl.TestFunction(M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b13645c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "g = dolfinx.fem.Function(V)\n",
    "g.interpolate(lambda x: np.sin(3 * x[0] + 1) * np.sin(3 * x[1] + 1))\n",
    "F = [(ufl.inner((1 + u**2) * ufl.grad(u), ufl.grad(v)) * ufl.dx\n",
    "      + ufl.inner(ufl.dot(ufl.grad(u), ufl.grad(u)) * u, v) * ufl.dx\n",
    "      + ufl.inner(l, v) * ufl.ds - ufl.inner(1, v) * ufl.dx),\n",
    "     ufl.inner(u, m) * ufl.ds - ufl.inner(g, m) * ufl.ds]\n",
    "J = [[ufl.derivative(F[0], u, du), ufl.derivative(F[0], l, dl)],\n",
    "     [ufl.derivative(F[1], u, du), ufl.derivative(F[1], l, dl)]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3053ba76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Class for interfacing with the SNES\n",
    "class NonlinearLagrangeMultplierBlockProblem:\n",
    "    \"\"\"Define a nonlinear problem, interfacing with SNES.\"\"\"\n",
    "\n",
    "    def __init__(  # type: ignore[no-any-unimported]\n",
    "        self, F: list[ufl.Form], J: list[list[ufl.Form]],\n",
    "        solutions: tuple[dolfinx.fem.Function, dolfinx.fem.Function],\n",
    "        bcs: list[dolfinx.fem.DirichletBC],\n",
    "        P: typing.Optional[list[list[ufl.Form]]] = None\n",
    "    ) -> None:\n",
    "        self._F = dolfinx.fem.form(F)\n",
    "        self._J = dolfinx.fem.form(J)\n",
    "        self._obj_vec = multiphenicsx.fem.petsc.create_vector_block(self._F, restriction)\n",
    "        self._solutions = solutions\n",
    "        self._bcs = bcs\n",
    "        self._P = P\n",
    "\n",
    "    def create_snes_solution(self) -> petsc4py.PETSc.Vec:  # type: ignore[no-any-unimported]\n",
    "        \"\"\"\n",
    "        Create a petsc4py.PETSc.Vec to be passed to petsc4py.PETSc.SNES.solve.\n",
    "\n",
    "        The returned vector will be initialized with the initial guesses provided in `self._solutions`,\n",
    "        properly stacked together and restricted in a single block vector.\n",
    "        \"\"\"\n",
    "        x = multiphenicsx.fem.petsc.create_vector_block(self._F, restriction=restriction)\n",
    "        with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(x, [V.dofmap, M.dofmap], restriction) as x_wrapper:\n",
    "            for x_wrapper_local, sub_solution in zip(x_wrapper, self._solutions):\n",
    "                with sub_solution.x.petsc_vec.localForm() as sub_solution_local:\n",
    "                    x_wrapper_local[:] = sub_solution_local\n",
    "        return x\n",
    "\n",
    "    def update_solutions(self, x: petsc4py.PETSc.Vec) -> None:  # type: ignore[no-any-unimported]\n",
    "        \"\"\"Update `self._solutions` with data in `x`.\"\"\"\n",
    "        x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "        with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(x, [V.dofmap, M.dofmap], restriction) as x_wrapper:\n",
    "            for x_wrapper_local, sub_solution in zip(x_wrapper, self._solutions):\n",
    "                with sub_solution.x.petsc_vec.localForm() as sub_solution_local:\n",
    "                    sub_solution_local[:] = x_wrapper_local\n",
    "\n",
    "    def obj(  # type: ignore[no-any-unimported]\n",
    "            self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec\n",
    "    ) -> np.float64:\n",
    "        \"\"\"Compute the norm of the residual.\"\"\"\n",
    "        self.F(snes, x, self._obj_vec)\n",
    "        return self._obj_vec.norm()  # type: ignore[no-any-return]\n",
    "\n",
    "    def F(  # type: ignore[no-any-unimported]\n",
    "        self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec, F_vec: petsc4py.PETSc.Vec\n",
    "    ) -> None:\n",
    "        \"\"\"Assemble the residual.\"\"\"\n",
    "        self.update_solutions(x)\n",
    "        with F_vec.localForm() as F_vec_local:\n",
    "            F_vec_local.set(0.0)\n",
    "        multiphenicsx.fem.petsc.assemble_vector_block(  # type: ignore[misc]\n",
    "            F_vec, self._F, self._J, self._bcs, x0=x, alpha=-1.0,\n",
    "            restriction=restriction, restriction_x0=restriction)\n",
    "\n",
    "    def J(  # type: ignore[no-any-unimported]\n",
    "        self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec, J_mat: petsc4py.PETSc.Mat,\n",
    "        P_mat: petsc4py.PETSc.Mat\n",
    "    ) -> None:\n",
    "        \"\"\"Assemble the jacobian.\"\"\"\n",
    "        J_mat.zeroEntries()\n",
    "        multiphenicsx.fem.petsc.assemble_matrix_block(\n",
    "            J_mat, self._J, self._bcs, diagonal=1.0,  # type: ignore[arg-type]\n",
    "            restriction=(restriction, restriction))\n",
    "        J_mat.assemble()\n",
    "        if self._P is not None:\n",
    "            P_mat.zeroEntries()\n",
    "            multiphenicsx.fem.petsc.assemble_matrix_block(\n",
    "                P_mat, self._P, self._bcs, diagonal=1.0,  # type: ignore[arg-type]\n",
    "                restriction=(restriction, restriction))\n",
    "            P_mat.assemble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ea532b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create problem\n",
    "problem = NonlinearLagrangeMultplierBlockProblem(F, J, (u, l), [])\n",
    "F_vec = multiphenicsx.fem.petsc.create_vector_block(problem._F, restriction=restriction)\n",
    "J_mat = multiphenicsx.fem.petsc.create_matrix_block(problem._J, restriction=(restriction, restriction))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36d66f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "snes = petsc4py.PETSc.SNES().create(mesh.comm)\n",
    "snes.setTolerances(max_it=20)\n",
    "snes.getKSP().setType(\"preonly\")\n",
    "snes.getKSP().getPC().setType(\"lu\")\n",
    "snes.getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "snes.setObjective(problem.obj)\n",
    "snes.setFunction(problem.F, F_vec)\n",
    "snes.setJacobian(problem.J, J=J_mat, P=None)\n",
    "snes.setMonitor(lambda _, it, residual: print(it, residual))\n",
    "solution = problem.create_snes_solution()\n",
    "snes.solve(None, solution)\n",
    "problem.update_solutions(solution)  # TODO can this be safely removed?\n",
    "solution.destroy()\n",
    "snes.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654d6229",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u, \"u\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e027a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(l, \"l\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cf083a1",
   "metadata": {},
   "source": [
    "### Strong imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f02f6c14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Class for interfacing with the SNES\n",
    "class NonlinearStrongBoundaryConditionsProblem:\n",
    "    \"\"\"Define a nonlinear problem, interfacing with SNES.\"\"\"\n",
    "\n",
    "    def __init__(  # type: ignore[no-any-unimported]\n",
    "        self, F: ufl.Form, J: ufl.Form, solution: dolfinx.fem.Function,\n",
    "        bcs: list[dolfinx.fem.DirichletBC], P: typing.Optional[ufl.Form] = None\n",
    "    ) -> None:\n",
    "        self._F = dolfinx.fem.form(F)\n",
    "        self._J = dolfinx.fem.form(J)\n",
    "        self._obj_vec = dolfinx.fem.petsc.create_vector(self._F)\n",
    "        self._solution = solution\n",
    "        self._bcs = bcs\n",
    "        self._P = P\n",
    "\n",
    "    def create_snes_solution(self) -> petsc4py.PETSc.Vec:  # type: ignore[no-any-unimported]\n",
    "        \"\"\"\n",
    "        Create a petsc4py.PETSc.Vec to be passed to petsc4py.PETSc.SNES.solve.\n",
    "\n",
    "        The returned vector will be initialized with the initial guess provided in `self._solution`.\n",
    "        \"\"\"\n",
    "        x = self._solution.x.petsc_vec.copy()\n",
    "        with x.localForm() as _x, self._solution.x.petsc_vec.localForm() as _solution:\n",
    "            _x[:] = _solution\n",
    "        return x\n",
    "\n",
    "    def update_solution(self, x: petsc4py.PETSc.Vec) -> None:  # type: ignore[no-any-unimported]\n",
    "        \"\"\"Update `self._solution` with data in `x`.\"\"\"\n",
    "        x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "        with x.localForm() as _x, self._solution.x.petsc_vec.localForm() as _solution:\n",
    "            _solution[:] = _x\n",
    "\n",
    "    def obj(  # type: ignore[no-any-unimported]\n",
    "        self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec\n",
    "    ) -> np.float64:\n",
    "        \"\"\"Compute the norm of the residual.\"\"\"\n",
    "        self.F(snes, x, self._obj_vec)\n",
    "        return self._obj_vec.norm()  # type: ignore[no-any-return]\n",
    "\n",
    "    def F(  # type: ignore[no-any-unimported]\n",
    "        self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec, F_vec: petsc4py.PETSc.Vec\n",
    "    ) -> None:\n",
    "        \"\"\"Assemble the residual.\"\"\"\n",
    "        self.update_solution(x)\n",
    "        with F_vec.localForm() as F_vec_local:\n",
    "            F_vec_local.set(0.0)\n",
    "        dolfinx.fem.petsc.assemble_vector(F_vec, self._F)\n",
    "        dolfinx.fem.petsc.apply_lifting(F_vec, [self._J], [self._bcs], x0=[x], alpha=-1.0)\n",
    "        F_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)\n",
    "        dolfinx.fem.petsc.set_bc(F_vec, self._bcs, x, -1.0)\n",
    "\n",
    "    def J(  # type: ignore[no-any-unimported]\n",
    "        self, snes: petsc4py.PETSc.SNES, x: petsc4py.PETSc.Vec, J_mat: petsc4py.PETSc.Mat,\n",
    "        P_mat: petsc4py.PETSc.Mat\n",
    "    ) -> None:\n",
    "        \"\"\"Assemble the jacobian.\"\"\"\n",
    "        J_mat.zeroEntries()\n",
    "        dolfinx.fem.petsc.assemble_matrix(  # type: ignore[misc]\n",
    "            J_mat, self._J, self._bcs, diagonal=1.0)  # type: ignore[arg-type]\n",
    "        J_mat.assemble()\n",
    "        if self._P is not None:\n",
    "            P_mat.zeroEntries()\n",
    "            dolfinx.fem.petsc.assemble_matrix(  # type: ignore[misc]\n",
    "                P_mat, self._P, self._bcs, diagonal=1.0)  # type: ignore[arg-type]\n",
    "            P_mat.assemble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ce9932",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "u_ex = dolfinx.fem.Function(V)\n",
    "F_ex = ufl.replace(F[0], {u: u_ex, l: 0})\n",
    "J_ex = ufl.derivative(F_ex, u_ex, du)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e984b5df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dirichlet BC object on Gamma\n",
    "dofs_V_Gamma = dolfinx.fem.locate_dofs_topological(V, boundaries.dim, facets_Gamma)\n",
    "bc_ex = [dolfinx.fem.dirichletbc(g, dofs_V_Gamma)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b83578e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create problem\n",
    "problem_ex = NonlinearStrongBoundaryConditionsProblem(F_ex, J_ex, u_ex, bc_ex)\n",
    "F_ex_vec = dolfinx.fem.petsc.create_vector(problem_ex._F)\n",
    "J_ex_mat = dolfinx.fem.petsc.create_matrix(problem_ex._J)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555e631e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "snes = petsc4py.PETSc.SNES().create(mesh.comm)\n",
    "snes.setTolerances(max_it=20)\n",
    "snes.getKSP().setType(\"preonly\")\n",
    "snes.getKSP().getPC().setType(\"lu\")\n",
    "snes.getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "snes.setObjective(problem_ex.obj)\n",
    "snes.setFunction(problem_ex.F, F_ex_vec)\n",
    "snes.setJacobian(problem_ex.J, J=J_ex_mat, P=None)\n",
    "snes.setMonitor(lambda _, it, residual: print(it, residual))\n",
    "solution_ex = problem_ex.create_snes_solution()\n",
    "snes.solve(None, solution_ex)\n",
    "problem_ex.update_solution(solution_ex)  # TODO can this be safely removed?\n",
    "solution_ex.destroy()\n",
    "snes.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3909209d",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u_ex, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9d09459",
   "metadata": {},
   "source": [
    "### Comparison and error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42bf8f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_ex_norm = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(ufl.grad(u_ex), ufl.grad(u_ex)) * ufl.dx)),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "err_norm = np.sqrt(mesh.comm.allreduce(\n",
    "    dolfinx.fem.assemble_scalar(dolfinx.fem.form(ufl.inner(ufl.grad(u_ex - u), ufl.grad(u_ex - u)) * ufl.dx)),\n",
    "    op=mpi4py.MPI.SUM))\n",
    "print(\"Relative error is equal to\", err_norm / u_ex_norm)\n",
    "assert np.isclose(err_norm / u_ex_norm, 0., atol=1.e-9)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
